package com.khalil.kkadmin.common.util;

import com.khalil.kkadmin.infrastructure.jpa.domain.TreeEntityBase;

import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.function.Predicate;
import java.util.stream.Collectors;

public class TreeUtil {

    public static <T extends TreeEntityBase<T, R>, R> List<T> buildTree(List<T> treeNodes, Predicate<T> isRoot) {
        Map<R, List<T>> treeNodeMap = treeNodes.stream().filter(i -> !isRoot.test(i)).collect(Collectors.groupingBy(TreeEntityBase::getParentId));

        for (T treeNode : treeNodes) {
            List<T> children = treeNodeMap.get(treeNode.getId());
            treeNode.setChildren(children);
        }

        return treeNodes.stream().filter(isRoot).collect(Collectors.toList());
    }

    public static <T extends TreeEntityBase<T, R>, R> void sortTree(List<T> treeNodes, Comparator<T> comparator) {
        if (treeNodes == null || treeNodes.isEmpty()) {
            return;
        }
        treeNodes.sort(comparator);
        for (T treeNode : treeNodes) {
            sortTree(treeNode.getChildren(), comparator);
        }
    }
}
